
import sys
import annotation

from theconfig import *

def buildSureSelectExonLookup(filename="/exds/bowtie/references/SureSelect-genes.txt"):
    lookup = {}
    f = open(filename, "r")
    for line in f:
        if line.startswith("#"): continue
        cols = line.strip().split("\t")
        if not lookup.has_key(cols[1]): lookup[cols[1]] = []
        # (exonStart, exonEnd, exonNum, geneName)
        exonStarts = cols[8].split(",")
        exonEnds = cols[9].split(",")
        for exonNum in xrange(int(cols[7])):
            lookup[cols[1]].append((int(exonStarts[exonNum]), int(exonEnds[exonNum])+1, exonNum+1, cols[0]))
    f.close()
    for key in lookup:
        lookup[key].sort()
        print >> sys.stderr, "# %d entries in SureSelect lookup table for %s" % (len(lookup[key]), key)
    return lookup

def getSureSelectExon(chr, pos, lookup):
    lookupChr = lookup[chr]
    for (spos, epos, exonNum, geneName) in lookupChr:
        if pos>=long(spos) and pos<=long(epos): return (geneName, exonNum)
    return ("None", 0)

def getCodons(positions, src, chromosomes=None):
    if src=="UCSC":
        codonInfoPattern = CODON_INFO_PATH + "/chr%s.ucsc-codon-info.txt"
        posShift = +1
    elif src=="CCDS":
        codonInfoPattern = CCDS_CODON_INFO_PATH + "/chr%s.ccds-codon-info.txt"
        posShift = 0
    else:
        raise Exception("src must be either UCSC or CCDS")
    ret = positions
    for key in ret.keys():
        ret[key] = []
    print >> sys.stderr, "# Getting %s codons for %d positions..." % (src, len(ret))
    totalHits = 0
    for x in range(1,23)+["X","Y"]:
        filename = codonInfoPattern % x
        if chromosomes and not chromosomes.has_key("chr%s" % x):
            print >> sys.stderr, "# ... skipping %s ..." % filename
            continue
        print >> sys.stderr, "# ... parsing %s ..." % filename
        hits = 0
        f = open(filename, "r")
        for line in f:
            cols = line.strip().split("\t")
            (name, chrom, strand, pos, codonNum, codon, frame) = tuple(cols)
            key = (chrom, long(pos) + posShift)
            if ret.has_key(key):
                hits += 1
                ret[key].append((name, strand, codonNum, codon, frame))
        f.close()
        print >> sys.stderr, "# ... %d hits." % hits
        totalHits += hits
    print >> sys.stderr, "# Total hits = ", totalHits
    return ret

def annotateSnps(file, shensnps, dbsnps, src="CCDS"):
    codonTable = annotation.readCodonTable()
    blosum = annotation.readBlosumMatrix(codonTable)

    snps = {}  # entries like "chr6_234524_A_T" --> tab-separated columns containing SNP id
    positions = {}  # entries like ("chr1", 1414) --> True
    rowContents = {}   # entries like "chr6_234524_A_T" --> tab-separated columns with add'l SNP info
    chromosomes = {}
    f = open(file)
    linenum = 0
    for line in f:
        linenum += 1
        if not line or line.startswith("#"): continue
        cols = line.strip().split("\t")
        if len(cols)>0 and cols[0].startswith("chr") and len(cols[0].split(":"))==2:
            # first column is chr:pos, second column is old base, third column is new base
            (chr, pos) = tuple(cols[0].split(":"))
            snpId = "%s_%s_%s_%s %d" % (chr, pos, cols[1], cols[2], linenum)
            if ("-" in pos) or ("random" in chr):
                print >> sys.stderr, "# Skipping line: ", line
                continue
            positions[(chr, long(pos))] = True
            snps[snpId] = "%s:%s\t%s\t%s" % (chr, pos, cols[1], cols[2])
            rowContents[snpId] = "\t".join(cols[3:])
        elif len(cols)==13 or len(cols)==16 or len(cols)==17 or len(cols)==18:
            # G 19 0.7600 chr1 4793 True A 30 0.764474 20 T 6/0/24/0 6/0/19/0
            snpId = "%s_%s_%s_%s %d" % (cols[3], cols[4], cols[6], cols[0], linenum)
            snps[snpId] = line.strip()
            positions[(cols[3], long(cols[4]))] = True
            rowContents[snpId] = ""
        elif len(cols)==22:
            # reading from snp-results.txt file:
            if cols[0]=="Sample": continue
            snps["chr%s_%s_%s_%s %d" % (cols[2], cols[3], cols[4], cols[5], linenum)] = line.strip()
            positions[("chr%s" % cols[2], long(cols[3]))] = True
        elif len(cols)>6:
            # arbitrary columns where 2nd column is SNP id (dropping first column)
            snpId = "%s %d" % (cols[1], linenum)
            snps[snpId] = "\t".join(cols[1].replace("chr", "").split("_"))
            fields = cols[1].split("_")
            positions[(fields[0], long(fields[1]))] = True
            rowContents[snpId] = "\t".join(cols[2:])
        else:
            snps["%s_%s_%s_%s %d" % (cols[0], cols[1], cols[2], cols[3], linenum)] = "%s\t%s\t%s\t%s" % (cols[0].replace("chr", ""), cols[1], cols[2], cols[3])
            positions[(cols[0], long(cols[1]))] = True
        if 0==len(positions)%10000:
            print >> sys.stderr, "# Read %d positions (%d snps) from file %s" % (len(positions), len(snps), file)
    for snpId in snps:
        chr = snpId.split("_")[0]
        chromosomes[chr] = True
    print >> sys.stderr, "# Read %d total positions in %d chromosomes (%d snps) from file %s" % (len(positions), len(chromosomes.keys()), len(snps), file)

    geneLookup = annotation.GeneInfoTable(src)

    annots = getCodons(positions, src, chromosomes)

    dbsnpFile = DBSNP_FILE.replace("snp131", "snp132")
    dbsnpNoString = "dbsnp132-no"
    
    ssExonLookup = buildSureSelectExonLookup()
    if not dbsnps:
        annotation.readDbsnpSnps(dbsnpFile, dbsnps, chromosomes)

    for snp in snps:
        # snpId = snp
        snpComponents = snp.split(" ")
        snpId = snpComponents[0]
        fields = snpId.split("_")
        (chr, pos, oldBase, newBase) = (fields[0], long(fields[1]), fields[2], fields[3])
        dbtxt = annotation.getDbsnpByMutation(dbsnps, chr, pos, oldBase, newBase) or dbsnpNoString
        shentxt = "shendure-snp-no"
        if shensnps.has_key(("%s:%d" % (chr, pos), oldBase, newBase)):
            shentxt = "shendure-snp-yes"
        annot = annots[(chr, pos)]
        ucscNames = geneLookup.getGeneIds(chr, pos) # list of geneIds
        symbol = geneLookup.getGeneSymbol(chr, pos) # first associated gene name, or "(None)"
        (ssGeneName, ssExonNum) = getSureSelectExon(chr, pos, ssExonLookup)
        ssExonTxt = "sureselect-no"
        if ssGeneName != "None":
            ssExonTxt = "sureselect-yes"
        ssEntry = "%s\t%s\t%d" % (ssExonTxt, ssGeneName, ssExonNum)
        foundAnnotation = False
        annotatedNames = {}
        for (ucscName, strand, codonNum, codon, frame) in annot:
            annotatedNames[ucscName] = True
        if len(ucscNames)==0:
            # if no associated transcripts at all, print a single row
            desc = "non-coding-unannotated"
            print "%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s" % (symbol, snpId, snps[snp], rowContents[snp], dbtxt, shentxt, "non-coding-unannotated", "\t".join(["NA"]*9), ssEntry)
            continue
        noncodingresults = {}
        for ucscName in ucscNames:
            # there are associated transcripts, combine those with same description of location
            if annotatedNames.has_key(ucscName):
                continue
            (strand, desc) = geneLookup.getNonCodingLocationDesc(chr, pos, ucscName)
            key = (strand, desc)
            if noncodingresults.has_key(key): noncodingresults[key] += ","
            noncodingresults[key] = noncodingresults.get(key, "") + ucscName
            if desc!="non-coding-gene":
                foundAnnotation = True
        results = {}
        codonNums = {}
        for (ucscName, strand, codonNum, codon, frame) in annot:
            foundAnnotation = True
            key = (strand, codon, int(frame))
            if results.has_key(key): results[key] += ","
            results[key] = results.get(key, "") + ucscName
            if codonNums.has_key(key): codonNums[key] += ","
            codonNums[key] = codonNums.get(key, "") + str(codonNum)
        for (strand, codon, frame) in results:
            if codon[frame] != annotation.toStrand(strand, oldBase):
                print >> sys.stderr, "WARNING: inconsistent old base at " + snp + " ... skipping this variant."
                continue
            oldCodon = ""
            newCodon = ""
            for i in [0,1,2]:
                if i!=frame:
                    oldCodon += codon[i]
                    newCodon += codon[i]
                else:
                    oldCodon += annotation.toStrand(strand, oldBase)
                    newCodon += annotation.toStrand(strand, newBase)
            oldAA = codonTable[oldCodon.upper()]
            newAA = codonTable[newCodon.upper()]
            oldAmino = oldAA[0]
            newAmino = newAA[0]
            proteinChanges = ",".join(["p.%s%d%s" % (oldAA[1], int(codonNum)+1, newAA[1]) for codonNum in codonNums[(strand, codon, frame)].split(",")])
            type = "syn"
            if oldAmino != newAmino:
                type = "non-syn"
            print "%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%d\t%s\t%s" % (symbol, snpId, snps[snp], rowContents[snp], dbtxt, shentxt, type, strand, frame+1, codon, newCodon, oldAmino, newAmino, proteinChanges, blosum[oldAmino][newAmino], results[(strand, codon, frame)], ssEntry)
        for (strand, desc) in noncodingresults:
            if not foundAnnotation or desc!="non-coding-gene":
                print "%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s" % (symbol, snpId, snps[snp], rowContents[snp], dbtxt, shentxt, desc, strand, "NA", "NA", "NA", "NA", "NA", "NA", "NA", noncodingresults[(strand, desc)], ssEntry)
    f.close()

if __name__=="__main__":
    if len(sys.argv)<=2 or sys.argv[1] not in ["UCSC", "CCDS"]:
        print >> sys.stderr, "Usage: python %s SRC=UCSC/CCDS [-D DIR] SNP-COUNTS" % sys.argv[0]
        sys.exit(1)
    src = sys.argv[1]
    sys.argv = sys.argv[1:]
    shensnps = annotation.readDbsnpSnps(SHENDURE_SNP_FILE)
    dbsnps = {}  # readDbsnpSnps()
    # dbsnps = None
    if sys.argv[1] == "-D":
        dir = sys.argv[2]
        for file in sys.argv[3:]:
            print >> sys.stderr, "# Annotating file ", file
            filehead = file.split("/")[-1].split(".")[0]
            tmpout = sys.stdout
            sys.stdout = open(dir + "/" + filehead + ".annotated.txt", "w")
            annotateSnps(file, shensnps, dbsnps, src)
            sys.stdout.close()
            sys.stdout = tmpout
    else:
        annotateSnps(sys.argv[1], shensnps, dbsnps, src)

